from langgraph.graph import StateGraph, END
from langgraph.prebuilt import ToolNode
from ..strategy_agent.backtest_engine.event_processor import (
    DataHandler,
    SignalGenerator,
    PortfolioManager,
    RiskChecker
)
import dask.dataframe as dd

class BacktestWorkflow:
    """事件驱动回测工作流"""
    
    def __init__(self):
        self.graph = StateGraph(dict)
        self._build_workflow()
        
    def _build_workflow(self):
        # 定义节点
        self.graph.add_node("data_loading", self.load_data)
        self.graph.add_node("signal_generation", self.generate_signals)
        self.graph.add_node("portfolio_management", self.manage_portfolio)
        self.graph.add_node("risk_check", self.check_risk)
        
        # 建立流程
        self.graph.set_entry_point("data_loading")
        self.graph.add_edge("data_loading", "signal_generation")
        self.graph.add_edge("signal_generation", "portfolio_management")
        self.graph.add_edge("portfolio_management", "risk_check")
        self.graph.add_edge("risk_check", END)
        
    async def load_data(self, state):
        """并行加载市场数据"""
        state['data'] = dd.read_parquet(
            f"data/{state['symbol']}/*.parquet",
            columns=['open', 'high', 'low', 'close', 'volume']
        )
        return state
    
    async def generate_signals(self, state):
        """分布式信号生成"""
        strategy = state['strategy']
        state['signals'] = state['data'].map_partitions(
            lambda df: df.apply(strategy.generate_signal, axis=1),
            meta=('signal', 'int')
        )
        return state
    
    async def manage_portfolio(self, state):
        """投资组合管理"""
        pm = PortfolioManager(state['capital'])
        state['positions'] = pm.execute_trades(
            state['signals'].compute(scheduler='threads')
        )
        return state
    
    async def check_risk(self, state):
        """实时风控检查"""
        risk_engine = RiskChecker()
        if not risk_engine.validate(state['positions']):
            raise RiskViolationError("风控规则触发")
        return state